import os
import re
from datetime import datetime
import xarray as xr
import numpy as np
# 运行目录需要放在根目录
def convert_nc_format(input_path, output_path, time_steps=2):
    """
    读取 NetCDF 文件，重命名变量，降采样，添加时间维度，并保存到新文件。
    """
    ds = xr.open_dataset(input_path)

    # 变量重命名映射
    rename_map = {
        "temperature_2m": "2m_temperature",
        "u_component_of_wind_10m": "10m_u_component_of_wind",
        "v_component_of_wind_10m": "10m_v_component_of_wind",
    }

    ds = ds.rename(rename_map)
    variables_to_keep = list(rename_map.values())
    ds = ds[variables_to_keep]

    # 降采样参数
    lat_size = ds.dims["latitude"]
    lon_size = ds.dims["longitude"]
    target_height = 60
    target_width = 120
    factor_lat = lat_size // target_height
    factor_lon = lon_size // target_width

    def downsample(da, factor_lat, factor_lon):
        lat_cut = (da.sizes["latitude"] // factor_lat) * factor_lat
        lon_cut = (da.sizes["longitude"] // factor_lon) * factor_lon
        da = da.isel(latitude=slice(0, lat_cut), longitude=slice(0, lon_cut))
        da = da.coarsen(latitude=factor_lat, longitude=factor_lon, boundary="trim").mean()
        return da

    def add_time_dim(da, time_steps):
        stacked = np.stack([da.values] * time_steps, axis=0)
        return xr.DataArray(
            data=stacked,
            dims=("time", "height", "width"),
            coords={
                "time": np.arange(time_steps),
                "height": np.arange(stacked.shape[1]),
                "width": np.arange(stacked.shape[2]),
            },
            name=da.name,
        )

    new_vars = {}
    for var in variables_to_keep:
        da = ds[var]
        da_down = downsample(da, factor_lat, factor_lon)
        da_time = add_time_dim(da_down, time_steps)
        new_vars[var] = da_time

    new_ds = xr.Dataset(new_vars)
    new_ds.to_netcdf(output_path)
    ds.close()
    print(f"保存文件: {output_path}")

def batch_process_nc_files(input_dir, output_dir, start_str, end_str):
    """
    批量转换指定时间范围内的 NetCDF 文件。
    文件名格式必须为 output_surface_YYYY-MM-DD-HH-MM.nc
    """
    os.makedirs(output_dir, exist_ok=True)

    pattern = re.compile(r"output_surface_(\d{4}-\d{2}-\d{2}-\d{2}-\d{2})\.nc")
    start_time = datetime.strptime(start_str, "%Y-%m-%d %H:%M")
    end_time = datetime.strptime(end_str, "%Y-%m-%d %H:%M")

    files = []
    for f in os.listdir(input_dir):
        m = pattern.match(f)
        if m:
            file_time = datetime.strptime(m.group(1), "%Y-%m-%d-%H-%M")
            if start_time <= file_time <= end_time:
                files.append((file_time, f))

    files.sort(key=lambda x: x[0])

    for i, (file_time, filename) in enumerate(files):
        input_path = os.path.join(input_dir, filename)
        output_filename = f"pgpred_{i:03d}.nc"
        output_path = os.path.join(output_dir, output_filename)
        print(f"处理文件 {filename} → {output_filename}")
        convert_nc_format(input_path, output_path, time_steps=2)

if __name__ == "__main__":
    input_folder = "2023-01-16-00-00to2023-01-20-23-00"  # 修改为你的输入文件夹路径
    output_folder = "panguresults"  # 输出文件夹
    start_time_str = "2023-01-16 00:00"
    end_time_str = "2023-01-20 23:00"

    batch_process_nc_files(input_folder, output_folder, start_time_str, end_time_str)
